package wikiextract.nlp.trainingset.x;

import java.io.IOException;
import java.io.Reader;

import learning.crf.inference.IParseScorer;
import learning.crf.model.Model;
import learning.crf.training.BasicParseScorer;
import learning.crf.training.CRF;
import learning.crf.training.CRFParameters;
import learning.crf.training.CRF.DecodingResult;
import learning.data.Dataset;
import learning.data.SequenceDataReader;
import learning.data.document.SequenceDocument;
import wikiextract.nlp.trainingset.x.Definitions.Feature1;

public class CRFTrainer {
	static int REPETITIONS = 1;
	
	private Eval eval;
	private int[] tpFpFnTn;
	private CRFParameters params;
	
	public CRFTrainer() {
		eval = new Eval();
		tpFpFnTn = new int[4];
	}
	
	public void addCV(Reader trData, Reader trLabels, Reader trFeatures,
			Reader teData, Reader teLabels, Reader teFeatures) 
		throws IOException {
		boolean transitions = Settings.useFeature1(Feature1.TRANSITIONS);
		
		SequenceDataReader.learnTransitionFeatures = transitions;
		// load datasets
		Dataset<SequenceDocument> trainData = 
			SequenceDataReader.read(trData, trLabels, trFeatures, -1);
		Dataset<SequenceDocument> testData = 
			SequenceDataReader.read(teData, teLabels, teFeatures, -1); 
		//System.out.println(trainData.numDocs() + " train, " + testData.numDocs() + " test");

		for (int r=0; r < REPETITIONS; r++) {
			// running cross validation			
			learning.crf.training.CollinsTraining ct = new learning.crf.training.CollinsTraining
				(transitions);
		
			CRFParameters params = ct.train(trainData);
			test2(testData, params, eval);
			// switch to test2 if only testing sentence
		}
	
	}
	
	public void train(Reader trData, Reader trLabels, Reader trFeatures) 
		throws IOException {
		boolean transitions = Settings.useFeature1(Feature1.TRANSITIONS);
		SequenceDataReader.learnTransitionFeatures = transitions;
		// load datasets
		Dataset<SequenceDocument> trainData = 
			SequenceDataReader.read(trData, trLabels, trFeatures, -1);
		learning.crf.training.CollinsTraining ct = new learning.crf.training.CollinsTraining
			(transitions);
		params = ct.train(trainData);
	}

	public void test(Reader teData, Reader teLabels, Reader teFeatures)
	throws IOException {
		test(teData, teLabels, teFeatures, null);
	}

	public void test(Reader teData, Reader teLabels, Reader teFeatures, TestListener listener)
	throws IOException {
		IParseScorer parseScorer = new BasicParseScorer(BasicParseScorer.Type.ALL_TERMINALS);
		test(teData, teLabels, teFeatures, parseScorer, listener);
	}

	public void test(Reader teData, Reader teLabels, Reader teFeatures, 
			IParseScorer parseScorer, TestListener listener)
		throws IOException {
		boolean transitions = Settings.useFeature1(Feature1.TRANSITIONS);
		SequenceDataReader.learnTransitionFeatures = transitions;
		Dataset<SequenceDocument> testData = 
			SequenceDataReader.read(teData, teLabels, teFeatures, -1); 
		test2(testData, params, eval, parseScorer, listener, tpFpFnTn);		
	}
	
	public void resetTest() {
		eval = new Eval();
		tpFpFnTn = new int[4];
	}
	
	public double[] getCrossValidationResults()  {
		double precision = eval.avgPrecision();
		double recall = eval.avgRecall();
		double f1 = 2.0* (precision*recall) / (precision + recall);
		if (Double.isNaN(f1)) f1 = 0;
		
		//System.out.println("   CV AVG f1 " + f1 + ", precision " + eval.avgPrecision() + ", recall " + eval.avgRecall());
		return new double[] { f1, precision, recall };
	}

	public int[] getTpFpFnTn() {
		return tpFpFnTn;
	}
	
	public static void test(Dataset<SequenceDocument> testData, CRFParameters parameters, Eval eval) {
		test(testData, parameters, eval, null);
	}
	
	public static void test(Dataset<SequenceDocument> testData, CRFParameters parameters, Eval eval, TestListener listener) {
		parameters.model = new Model(testData.numLabels());

		BasicParseScorer allParseScorer = new BasicParseScorer(BasicParseScorer.Type.ALL_TERMINALS);
		allParseScorer.setDataset(testData);
		
		int truePos = 0; // articles (or words?)
		int falsePos = 0;
		int trueNeg = 0;
		int falseNeg = 0;
		int incorrectLabel = 0;
		
		String extracted = null;
		String labeled = null;
		int prevArticleId = -1;
		int articleId = -1;
		for (SequenceDocument doc : testData) {
			articleId = Integer.parseInt(doc.meta.get("articleId"));

			if (articleId != prevArticleId && prevArticleId != -1) {
				if (extracted == null && labeled != null) falseNeg++;
				if (extracted != null && labeled == null) falsePos++;
				if (extracted == null && labeled == null) trueNeg++;
				if (extracted != null && labeled != null &&
						extracted.equals(labeled)) truePos++;
				if (extracted != null && labeled != null &&
						!extracted.equals(labeled)) incorrectLabel++;
				/*				
				if (extracted != null || labeled != null) {
					System.out.println(" predicted: " + extracted);
					System.out.println(" actual:    " + labeled);
				}*/
				extracted = null;
				labeled = null;
			}
				
			// compute most likely label under current parameters
			DecodingResult predictedParse = CRF.decode(doc, allParseScorer, parameters);
			DecodingResult trueParse = new DecodingResult(doc.labels, 0);
			
			if (extracted == null)
				for (int i=0; i < predictedParse.labels.length; i++) {
					if (predictedParse.labels[i] == 1) {
						if (extracted == null) extracted = "";
						extracted += doc.tokens[i] + " ";
					}
				}
			
			if (labeled == null)
				for (int i=0; i < trueParse.labels.length; i++) {
					if (trueParse.labels[i] == 1) {
						if (labeled == null) labeled = "";
						labeled += doc.tokens[i] + " ";
					}
				}
			
			
			// update evaluator
			//evaluator.update(doc, predictedParse, trueParse);
			prevArticleId = articleId;
		}
		// last one
		if (prevArticleId != -1) {
			if (extracted == null && labeled != null) falseNeg++;
			if (extracted != null && labeled == null) falsePos++;
			if (extracted == null && labeled == null) trueNeg++;
			if (extracted != null && labeled != null &&
					extracted.equals(labeled)) truePos++;
			if (extracted != null && labeled != null &&
					!extracted.equals(labeled)) incorrectLabel++;
			//System.out.println("labeled " + labeled);
		}
		
		double precision = truePos / (double)(incorrectLabel + truePos + falsePos);
		double recall = truePos / (double)(truePos + falseNeg + incorrectLabel);
		if (testData.numDocs() == 0) System.out.println("warning: test data size = 0");
		if (Double.isNaN(precision)) precision = 1; // not
		if (Double.isNaN(recall)) recall = 1; // not
		eval.update(precision, recall);
		//System.out.println("    precision " + precision + ", recall " + recall);
	}

	public static void test2(Dataset<SequenceDocument> testData, CRFParameters parameters, Eval eval) {
		IParseScorer parseScorer = new BasicParseScorer(BasicParseScorer.Type.ALL_TERMINALS);
		test2(testData, parameters, eval, parseScorer, null, null);
	}

	public static void test2(Dataset<SequenceDocument> testData, CRFParameters parameters, Eval eval, 
			IParseScorer parseScorer, TestListener listener,int[] tpFpFnTn) {
		parameters.model = new Model(testData.numLabels());

		parseScorer.setDataset(testData);
		
		int truePos = 0; // articles (or words?)
		int falsePos = 0;
		int trueNeg = 0;
		int falseNeg = 0;
		
		for (SequenceDocument doc : testData) {
			// compute most likely label under current parameters
			DecodingResult predictedParse = CRF.decode(doc, parseScorer, parameters);
			DecodingResult trueParse = new DecodingResult(doc.labels, 0);

			for (int i=0; i < trueParse.labels.length; i++) {
				if (trueParse.labels[i] == 1 && predictedParse.labels[i] == 1) truePos++;
				if (trueParse.labels[i] == 1 && predictedParse.labels[i] == 0) falseNeg++;
				if (trueParse.labels[i] == 0 && predictedParse.labels[i] == 1) falsePos++;
				if (trueParse.labels[i] == 0 && predictedParse.labels[i] == 0) trueNeg++;
			}
			if (listener != null)
				listener.next(doc, predictedParse.labels);
		}
		if (tpFpFnTn != null) {
			tpFpFnTn[0] += truePos;
			tpFpFnTn[1] += falsePos;
			tpFpFnTn[2] += falseNeg;
			tpFpFnTn[3] += trueNeg;
		}
		
		double precision = truePos / (double)(truePos + falsePos);
		double recall = truePos / (double)(truePos + falseNeg);
		if (testData.numDocs() == 0) System.out.println("warning: test data size = 0");
		if (Double.isNaN(precision)) precision = 1; // not
		if (Double.isNaN(recall)) recall = 1; // not
		eval.update(precision, recall);		
	}
	
	
	static class Eval {
		int num = 0;
		double precision = 0;
		double recall = 0;
		
		public void update(double precision, double recall) {
			this.precision += precision;
			this.recall += recall;
			this.num++;
		}
		
		public double avgPrecision() {
			return precision / num;
		}
		
		public double avgRecall() {
			return recall / num;
		}
	}

	static interface TestListener {
		public void next(SequenceDocument doc, int[] predicted);
	}
}
	
